import os
import torch
from tqdm import tqdm
from eval_coco import create_batches_from_json

# Set the CUDA device to use
os.environ['CUDA_VISIBLE_DEVICES'] = '1'  # Replace '0' with the desired CUDA device index

def save_coco_features_and_ids(model, json_path, base_folder, batch_size=32, feature_dir="./features", chunk_size=1024):
    """
    Save features and IDs for COCO dataset for cross-modal retrieval.

    Args:
        model: The model used to extract features.
        json_path: Path to the COCO annotations JSON file.
        base_folder: Root folder containing the COCO dataset.
        batch_size: Number of samples per batch.
        feature_dir: Directory to save the features and IDs.
        chunk_size: Number of features to save per file to prevent memory overflow.
    """
    # Create batches for images and texts
    print("Creating batches for images...")
    image_batches, image_ids = create_batches_from_json(json_path, base_folder, batch_size, modality="image")

    print("Creating batches for texts...")
    text_batches, text_ids = create_batches_from_json(json_path, base_folder, batch_size, modality="text")

    # Ensure feature directory exists
    os.makedirs(feature_dir, exist_ok=True)

    # Save image features and IDs
    print("Extracting and saving image features...")
    image_feat_list = []
    chunk_index = 0
    for i, batch in enumerate(tqdm(image_batches, desc="Processing image batches")):
        features = model.emb_images(batch)
        image_feat_list.append(features)

        if len(image_feat_list) >= chunk_size:
            intermediate_features = torch.cat(image_feat_list, dim=0)
            # Print the current chunk index being saved
            print(f"Saving chunk {chunk_index} for image features...")
            torch.save(intermediate_features, os.path.join(feature_dir, f"image_features_chunk_{chunk_index}.pt"))
            torch.save(image_ids[chunk_index * chunk_size:(chunk_index + 1) * chunk_size], os.path.join(feature_dir, f"image_ids_chunk_{chunk_index}.pt"))
            image_feat_list.clear()
            chunk_index += 1

    if image_feat_list:
        intermediate_features = torch.cat(image_feat_list, dim=0)
        torch.save(intermediate_features, os.path.join(feature_dir, f"image_features_chunk_{chunk_index}.pt"))
        torch.save(image_ids[chunk_index * chunk_size:], os.path.join(feature_dir, f"image_ids_chunk_{chunk_index}.pt"))

    # Save text features and IDs
    print("Extracting and saving text features...")
    text_feat_list = []
    chunk_index = 0
    for i, batch in enumerate(tqdm(text_batches, desc="Processing text batches")):
        features = model.emb_texts(batch)
        text_feat_list.append(features)

        if len(text_feat_list) >= chunk_size:
            intermediate_features = torch.cat(text_feat_list, dim=0)
            # Print the current chunk index being saved for text features
            print(f"Saving chunk {chunk_index} for text features...")
            torch.save(intermediate_features, os.path.join(feature_dir, f"text_features_chunk_{chunk_index}.pt"))
            torch.save(text_ids[chunk_index * chunk_size:(chunk_index + 1) * chunk_size], os.path.join(feature_dir, f"text_ids_chunk_{chunk_index}.pt"))
            text_feat_list.clear()
            chunk_index += 1

    if text_feat_list:
        intermediate_features = torch.cat(text_feat_list, dim=0)
        torch.save(intermediate_features, os.path.join(feature_dir, f"text_features_chunk_{chunk_index}.pt"))
        torch.save(text_ids[chunk_index * chunk_size:], os.path.join(feature_dir, f"text_ids_chunk_{chunk_index}.pt"))

    # Print the number of features saved for each modality
    print("Number of image features saved:", len(image_ids))
    print("Number of text features saved:", len(text_ids))

if __name__ == "__main__":
    from omni_model.omni_space import OmniBind_Base

    # Load the model
    model = OmniBind_Base(pretrained=True).cuda().eval()

    # Define paths
    json_path = "/data/jzw/MSCOCO/annotations/captions_train2017.json"
    base_folder = "/data/jzw/MSCOCO/train2017"
    feature_dir = "./coco_train2017_features"

    # Save features and IDs
    save_coco_features_and_ids(model, json_path, base_folder, batch_size=32, feature_dir=feature_dir, chunk_size=1024)